from .base import pool, cates, BaseGrouping
from models import Sample
from typing import List, Set
from concurrent.futures import Future
from event_bus import EventBus
import os
import csv


class CSVGroupingEngine(BaseGrouping):
    task_record: List[Future] = []

    def __init__(self, export_location: str, category_file: str = "category", sample_file: str = "sample",
                 validation_file: str = "validation", resize_to: int = 0, padding: bool = False,
                 sample_type: str = "png", categories: Set[str] = None, group: int = 1, training: float = 0.,
                 validate: float = 0.):
        super().__init__(export_location, category_file, sample_file, validation_file,
                         resize_to, padding, sample_type, categories, group, training, validate)

    def calculate_works(self) -> int:
        total_samples = Sample.select() \
            .from_(Sample, cates) \
            .where(cates.c.value << self.categories) \
            .distinct() \
            .count()
        return total_samples

    def split_group(self):
        import random
        selected_samples = Sample.select(Sample) \
            .from_(Sample, cates) \
            .where(cates.c.value << self.categories) \
            .distinct()
        fingerprints = {s.target_fingerprint for s in selected_samples.objects()}
        group_size = len(fingerprints) // self.split_groups
        self.groups.clear()
        for i in range(self.split_groups):
            chooses = set(random.sample(fingerprints, group_size))
            fingerprints = fingerprints - chooses
            training_size = round(len(chooses) * self.training_ratio)
            training_group = set(random.sample(chooses, training_size))
            self.groups.append((training_group, chooses - training_group))

    def _write_csv_description(self, target_file: str, fingerprints: Set[str]):
        try:
            file_affix = "png" if self.sample_image_type.lower() == "png" else "jpg"
            with open(target_file, 'w') as desc:
                writer = csv.writer(desc)
                writer.writerow(["文件名", *self.categories])
                for fp in fingerprints:
                    rec = Sample.get(Sample.target_fingerprint == fp)
                    class_labels = [1 if c in rec.categories else 0
                                    for c in self.categories]
                    writer.writerow([f"{rec.target_fingerprint}.{file_affix}", *class_labels])
        except OSError as e:
            EventBus().post_topic(event=e, topic="export_failed")
            return

    def write_group_description(self, group: int):
        if group < len(self.groups):
            working_group = self.groups[group]
            try:
                group_location = os.path.join(self.export_location, f"group{group:03d}")
                if not os.path.exists(group_location):
                    os.makedirs(group_location)
                self.task_record.append(
                    pool.submit(
                        self._write_csv_description,
                        os.path.join(group_location, f"{self.sample_file}.csv"),
                        working_group[0]
                    )
                )
                self.task_record.append(
                    pool.submit(
                        self._write_csv_description,
                        os.path.join(group_location, f"{self.validation_file}.csv"),
                        working_group[1]
                    )
                )
            except OSError as e:
                EventBus().post_topic(event=e, topic="export_failed")
                return

    def generate_sample_file(self, group: int):
        if group < len(self.groups):
            working_group = self.groups[group]
            samples = working_group[0] | working_group[1]
            group_location = os.path.join(self.export_location, f"group{group:03d}")
            if not os.path.exists(group_location):
                os.makedirs(group_location)
            for fp in samples:
                sample = Sample.get(Sample.target_fingerprint == fp)
                self.task_record.append(
                    pool.submit(
                        self.save_sample_picture,
                        group_location,
                        sample
                    )
                )

    def generate_samples(self):
        for i in range(self.split_groups):
            self.write_group_description(i)
            self.generate_sample_file(i)

    def cancel(self):
        for task in self.task_record:
            task.cancel()
